import base64
from fastapi import FastAPI
from pydantic import BaseModel
import ddddocr
import uvicorn

ocr = ddddocr.DdddOcr()
app = FastAPI()


class Demo(BaseModel):
    image_base64: str


@app.get("/test")
async def root():
    return {"message": "Hello World"}


@app.post("/verification-code")
def parse_captcha(demo: Demo):
    print("parse captcha")
    print(demo.image_base64)
    image_base64 = demo.image_base64

    if not image_base64:
        return {"error": "缺少图片数据"}

    # 解码 Base64 数据
    try:
        if image_base64.startswith("data:image"):
            header, image_data = image_base64.split(",", 1)

        # 将 Base64 转换为字节流
        image_bytes = base64.b64decode(image_data)

        # 使用 ddddocr 解析图片
        result = ocr.classification(image_bytes)
        print(f"result: {result}")

        return {"captcha_value": result}
    except Exception as e:
        return {"error": str(e)}


if __name__ == "__main__":
    uvicorn.run("main:app", host="0.0.0.0", port=8000)
